rm(list=ls())
setwd("~/Desktop/covid2/final code and data/data analysis/full model/") 
source("functions_full.r")
library(splines)
library(CBPS)
library(zoo)



lag   <- 0    
iters <- 200000
burn  <- iters/5



## load and format data
load("../data.rdata") 
if (exists(".Random.seed")) { rm(.Random.seed) }
YnonCum <- matrix(NA,nrow=nrow(Y), ncol=ncol(Y)-1)
for (i in 1:nrow(Y)) {
  YnonCum[i,] <- diff(Y[i,])
}  
Y <- YnonCum
Y <- Y[,5:221]
Y[Y<0] <- 0
range(dateRange[6:222])
for (i in 1:length(allA)) { 
  if (!is.null(allA[[i]])) {
    allA[[i]] <- allA[[i]][,6:222]
  }
}
for (i in 1:length(allX)) { 
  if (!is.null(allX[[i]])) {
    allX[[i]] <- allX[[i]][,6:222] 
  }
}
N  <- N[,6:222]
nt <- length(6:222)
A  <- allA[[7]]
nT <- nt/7
weeklyAvg <- function(Win) {
  Wout <- matrix(NA, nrow=ns, ncol=nT)  
  for (i in 1:ns) { for (j in 1:nT) { Wout[i,j] <- mean(Win[i, (j-1)*7 + 1:7]) } }
  return(Wout)
}
A  <- weeklyAvg(A)
for (i in 1:length(allX)) { allX[[i]] <- weeklyAvg(allX[[i]]) }

Aoriginal <- A
Y  <- round(weeklyAvg(Y)*7)
N  <- weeklyAvg(N)

bw <- 0.001
p  <- 0.5

Abar <- Ybar <- Nbar <- matrix(0,ns,nT)
allXbar <- list()
for (i in 1:length(allX)) { allXbar[[i]] <- matrix(0,ns,nT) }
for (s in 1:ns){ for(t in 1:nT){
  Abar[s,t]    <- mean(A[adjs[[s]],t])
  for (i in 1:length(allX)) { allXbar[[i]][s,t] <- mean(allX[[i]][adjs[[s]],t]) }
}}
w  <- make.w(nT,lag,bw)  # compute weighted function for reporting lag
id <- 2:nT
B  <- bs(1:nT, degree=3) 
allX[[8]] <- log(allX[[8]])

df <- cbind(c(A[,id]),            # A_t
            # lagged A values:
            c( A[,id-1] ),        # A_t-1
            # time-varying covariates (current and lagged):
            c( allX[[1]][,id] ) , # meantemp_t
            c( allX[[1]][,id-1]), # meantemp_t-1
            c( allX[[2]][,id] ) , # hightemp_t
            c( allX[[2]][,id-1]), # hightemp_t-1
            c( allX[[3]][,id] ) , # lowtemp_t
            c( allX[[3]][,id-1]), # lowtemp_t-1
            c( allX[[4]][,id] ) , # relhumidity_t
            c( allX[[4]][,id-1]), # relhumidity_t-1
            c( allX[[5]][,id] ) , # dewpoint_t
            c( allX[[5]][,id-1]), # dewpoint_t-1
            # static covariates:
            c( allX[[6]][,id] ) , # pm25
            c( allX[[7]][,id] ) , # poverty
            c( allX[[8]][,id] ) , # popdens
            c( allX[[9]][,id] ) , # medhouseinc
            c( allX[[10]][,id] ), # population
            c( allX[[11]][,id] ), # beds
            c( allX[[12]][,id] ), # icubeds
            c( allX[[13]][,id] ), # medianage
            c( allX[[14]][,id] ), # foreignborn
            c( allX[[15]][,id] ), # healthemp
            c( allX[[16]][,id] ), # daysincefirstcaseinstate
            c( allX[[16]][,id]^2),# daysincefirstcaseinstate^2
            c(t(replicate(ns,1:nt))[,id]), # time (1:nt)
            c(t(replicate(ns,(1:nt)^2))[,id]), # time^2 (1:nt)
            c( t(replicate(ns,1:nt))[,id] * A[,id-1] ), # time*A_t-1
            c( replicate(nT,Aoriginal[,1])[,id] ),
            c( log(Y[,id-1]+1)-log(N[,1]+1) ))#,  # log(Y_t-1) - log(N_j +1)
df <- as.data.frame(df)
names(df) <- c("A_t","A_tm1","meantemp_t","meantemp_tm1","hightemp_t","hightemp_tm1",
               "lowtemp_t","lowtemp_tm1","relhumidity_t","relhumidity_tm1",
               "dewpoint_t","dewpoint_tm1","pm25","poverty","popdens","medhouseinic",
               "population","beds","icubeds","medianage","foreignborn","healthemp",
               "daysincefirstcaseinstate","daysincefirstcaseinstate2",
               "time","time2", "timeA_tm1", "baselineA", "logYtm1_m_logN")
summary(ols1 <- lm(A_t ~ ., data=df))

Ahat <- A
Ahat[,id]    <- ols1$fitted
Ahatbar      <- matrix(0,ns,nT)
for(s in 1:ns){for(t in 1:nT){ Ahatbar[s,t] <- mean(Ahat[adjs[[s]],t]) }}

cnt <- 0
X   <- list()
X[[cnt<-cnt+1]] <- matrix(1,ns,nT)
X[[cnt<-cnt+1]] <- A - Ahat      
X[[cnt<-cnt+1]] <- Abar - Ahatbar
X[[cnt<-cnt+1]] <- Ahat
X[[cnt<-cnt+1]] <- Ahatbar
X[[cnt<-cnt+1]] <- allX[[1]]   # meantemp_t
X[[cnt<-cnt+1]] <- allX[[4]]   # relhumidity_t
X[[cnt<-cnt+1]] <- allX[[5]]   # dewpoint_t
X[[cnt<-cnt+1]] <- allX[[6]]   # pm25
X[[cnt<-cnt+1]] <- allX[[7]]   # poverty
X[[cnt<-cnt+1]] <- allX[[8]]   # popdens
X[[cnt<-cnt+1]] <- allX[[9]]   # medhouseinic
X[[cnt<-cnt+1]] <- allX[[10]]  # population
X[[cnt<-cnt+1]] <- allX[[11]]  # beds
X[[cnt<-cnt+1]] <- allX[[12]]  # icubeds 
X[[cnt<-cnt+1]] <- allX[[13]]  # medianage
X[[cnt<-cnt+1]] <- allX[[14]]  # foreignborn
X[[cnt<-cnt+1]] <- allX[[15]]  # healthemp
X <- lapply(X,lag_mat,lag=lag,nt=nT)

names(X) <- c("intercept","A", "Abar","Ahat", "Ahatbar","meantemp","relhumidity",
"dewpoint","pm25","poverty","popdens","medhouseinc","population","beds",
"icubeds","medianage","foreignborn","healthemp")

fit <- MyCAR(Y,adjs,X=X,N=N,firstday=8,nugget=T,spatial=T,
             iters=iters,burn=burn,thin=1,update=100,
             showProgress=100)
saveRDS(fit, file="fit.rds")

